from dotenv import load_dotenv
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
import json
import asyncio
import os
import dashscope
import logging
from dashscope.audio.asr import *
from logger import logging

router = APIRouter(prefix="/api")

# 初始化API密钥
load_dotenv()
dashscope.api_key = os.getenv('API_KEY')
# Real-time speech recognition callback


class MyRecognitionCallback(RecognitionCallback):
    def __init__(self, tag, websocket, loop) -> None:
        super().__init__()
        self.tag = tag
        self.text = ''
        self.websocket = websocket
        self.loop = loop
        self.send_events = []

    def on_open(self) -> None:
        pass

    def on_complete(self) -> None:
        pass

    def on_error(self, result: RecognitionResult) -> None:
        logging.error(f"[asr]回调发送错误: {result}")

    async def send_asr_result(self, message: str,
                              event: asyncio.Event) -> None:
        data = json.loads(message)
        if (data.get('is_end', False)):
            logging.info(f"[asr]识别结果: {data['text']}")
            try:
                await self.websocket.send_text(message)
            except Exception as e:
                logging.error(f'[asr]发送识别结果,客户端连接关闭: {e}')
        event.set()

    def on_event(self, result: RecognitionResult) -> None:
        sentence = result.get_sentence()
        if 'text' in sentence:
            is_end = False
            if RecognitionResult.is_sentence_end(sentence):
                is_end = True
            msg = {'text': sentence['text'], 'is_end': is_end}
            event = asyncio.Event()
            self.send_events.append(event)
            self.loop.call_soon_threadsafe(
                asyncio.create_task,
                self.send_asr_result(json.dumps(msg), event))

    def on_close(self) -> None:
        pass


@router.websocket("/ws/asr")
async def websocket_asr(websocket: WebSocket):
    """WebSocket语音识别端点"""
    await websocket.accept()
    logging.info("ASR WebSocket连接已建立")

    try:
        loop = asyncio.get_event_loop()
        callback = MyRecognitionCallback('process0', websocket, loop)
        recognition = Recognition(
            model='paraformer-realtime-v2',
            format='pcm',
            sample_rate=16000,  # supported 8000、16000
            semantic_punctuation_enabled=False,
            callback=callback)

        recognition.start()
        logging.info("语音识别服务已启动")

        try:
            while True:
                data = await websocket.receive_bytes()
                if isinstance(data, bytes):
                    recognition.send_audio_frame(data)

        except WebSocketDisconnect:
            logging.info("ASR WebSocket连接已断开")
        except Exception as e:
            logging.error(f'[asr]服务异常: {e}')
            try:
                await websocket.send_text(json.dumps({
                    'type': 'error',
                    'message': f'ASR服务器错误: {str(e)}'
                }))
            except:
                pass
        finally:
            recognition.stop()
            logging.info("语音识别服务已停止")

            # wait until all asr results are sent
            for event in callback.send_events:
                await event.wait()

    except Exception as e:
        logging.error(f"ASR WebSocket错误: {e}")
        try:
            await websocket.send_text(json.dumps({
                'type': 'error',
                'message': f'ASR服务器错误: {str(e)}'
            }))
        except:
            pass
